#import "LazyRegionManager.h"

#import <string.h>

#define 	VM_PROT_VOLATILE	((vm_prot_t) 0x08)	// not cacheable
#define	 	VM_PROT_SPARSE		((vm_prot_t) 0x10)	// sparse addr space

@implementation LazyRegionManager

-(BOOL)readInReloc: (Reloc *)reloc
{
    kern_return_t ret;
    
    if (!reloc->rFlags.readIn) {
	ret = vm_allocate(task_self(), &reloc->data, reloc->size, YES);
	if (ret == KERN_SUCCESS) {
	    reloc->maxData = reloc->data + reloc->size;
	    reloc->displacement = reloc->data - reloc->address;
	    reloc->rFlags.readIn = YES;
	}
    }
    return reloc->rFlags.readIn ? YES : NO;
}

-(int)pageFor: (const void *)pointer inRegion: (Region *)region
{
    return (((unsigned int)pointer - region->reloc.address) / pageSize);
}

-(BOOL)readInDataFor: (const void *)pointer into: (void *)data withSize: (int)size
{
    unsigned int error = 0, nData;
    pointer_t readPage;
    
    error = vm_read(task, (vm_address_t)pointer, size, &readPage, &nData);
    if (!error) {
	bcopy((void *)readPage, data, size);
	vm_deallocate(task_self(), readPage, nData);
    } else
	[self isTask];
    return error ? NO : YES;
}

-(BOOL)readInPage: (unsigned char *)page
              for: (void *)pointer
             into: (void *)data
{
    unsigned int error = 0, nData;
    pointer_t readPage;
    
    if (!(*page & PAGEREADIN)) {
//	printf("Reading in page at %p.\n", pointer);
	error = vm_read(task, (vm_address_t)pointer, pageSize, &readPage, 
	                &nData);
	if (!error) {
	    bcopy((void *)readPage, data, pageSize);
	    vm_deallocate(task_self(), readPage, nData);
	    if (!(*page & VM_PROT_VOLATILE))
		*page |= PAGEREADIN;
	} else
	    [self isTask];
    }
    return error ? NO : YES;
}

-(BOOL)readInPagesFor: (const void *)pointer withSize: (int)size inRegion: (Region *)region
{
    int nPage, offset;
    unsigned char *page, *endPage;
    void *data, *address;
    int ok = YES;
    
//    printf("Getting data at: %p for %d bytes.\n", pointer, size);
    
    nPage = [self pageFor: pointer inRegion: region];
    offset = nPage * pageSize;
    page = region->pages + nPage;
    if (*page & VM_PROT_SPARSE)
	ok = [self readInDataFor: pointer
	                    into: (void *)pointer + region->reloc.displacement 
			withSize: size];
    else {
	endPage = region->pages
		    + [self pageFor: pointer + size inRegion: region];
	data = (void *)(region->reloc.data + offset);
	address = (void *)(region->reloc.address + offset);
	while (ok && page <= endPage) {
	    ok = [self readInPage: page for: address into: data];
	    address += pageSize;
	    data += pageSize;
	    page++;
	}
    }
    return ok;
}

-(void *)pointerFor: (const void *)pointer withSize: (int)size
{
    Region *region;
    pointer_t newPointer;
    if (pointer && (region = (Region *)[self relocFor: pointer])) {
	newPointer = (pointer_t)pointer + region->reloc.displacement;
	if (((newPointer + size) <= region->reloc.maxData)
	    && [self readInPagesFor: pointer withSize: size inRegion: region])
	    return (void *)newPointer;
	else
	    return NULL;
    } else
	return NULL;
}

-(int)getDataAt: (const void *)start for: (int)numBytes into: (void *)data
{
    Reloc *reloc = [self relocFor: start];
    int numBytesInRegion;
    
//    printf("Getting data at: %p for %d bytes.\n", start, numBytes);
    
    if (reloc) {
        numBytesInRegion = reloc->maxAddress - (int)start;
        if (numBytes > numBytesInRegion)
	    numBytes = numBytesInRegion;
	if ([self readInPagesFor: start
	                withSize: numBytes
		        inRegion: (Region *)reloc]) {
	    memcpy(data,
		    (void *)(reloc->data
		             + ((pointer_t)start - reloc->address)),
		    numBytes);
	    return numBytes;
	} else
	    return 0;
    } else
	return 0;
}

-(int)putDataAt: (void *)start for: (int)numBytes from: (const void *)data
{
    int numBytesInRegion;
    Reloc *reloc = [self relocFor: start];
 
//    printf("Putting data at: %p for %d bytes.\n", start, numBytes);

    if (reloc) {
        numBytesInRegion = reloc->maxAddress - (int)start;
        if (numBytes > numBytesInRegion)
	    numBytes = numBytesInRegion;
	if ([self readInPagesFor: start
	                withSize: numBytes
		        inRegion: (Region *)reloc]) {
	    memcpy((void *)(reloc->data + ((pointer_t)start - reloc->address)),
		   data,
		   numBytes);
	    return [self writeDataAt: start for: numBytes
	                       reloc: (Region *)reloc];
	} else
	    return 0;
    } else
	return 0;
}

-(char *)pointerForString: (const char *)pointer isNullTerminated: (BOOL *)isNT
{
    char *retPointer;
    Region *region;
    char *string;
    int nPage, offset;
    unsigned char *page;
    void *data, *address;
    BOOL ok;
    
    if (pointer && (region = (Region *)[self relocFor: pointer])) {
    	retPointer = (char *)((pointer_t)pointer + region->reloc.displacement);
	nPage = [self pageFor: pointer inRegion: region];
	page = region->pages + nPage;
	offset = nPage * pageSize;
	data = (void *)(region->reloc.data + offset);
	address = (void *)(region->reloc.address + offset);
	string = retPointer;
	do {
	    ok = [self readInPage: page for: address into: data];
	    if (ok) {
		page++;
		address += pageSize;
		data += pageSize;
		while (*string && (string < (char *)data))
		    string++;
	    }
	} while (ok && *string && (data < (void *)region->reloc.maxData));
	*isNT = *string ? NO : YES;
	return retPointer;
    } else{
	*isNT = NO;
	return NULL;
    }
}

-(int)writeDataAt: (const void *)start for: (int)numBytes reloc: (Region *)region
{
    char *page = pages + [self pageFor: start inRegion: region];
    kern_return_t ret;
    
    if (*page & VM_PROT_SPARSE) {
	ret = vm_write(task,
	               (vm_address_t)start,
	               (pointer_t)(start + region->reloc.displacement),
		       numBytes);
	if (ret == KERN_SUCCESS)
	    return numBytes;
	else {
	    [self isTask];
	    return 0;
	}
    } else
	return [super writeDataAt: start for: numBytes reloc: region];
}

@end